from genericpath import exists
from sklearn.metrics import confusion_matrix, precision_recall_curve, f1_score
import numpy as np
import os


def cal_confusion_matrix(y_true, y_pred_no_thresh, thresh, img_path_list):
    pred_thresh = []
    false_n = []
    false_p = []
    for i in range(len(y_pred_no_thresh)):
        if y_pred_no_thresh[i] > thresh:
            pred_thresh.append(1)
            if y_true[i] == 0:
                false_p.append(img_path_list[i])   # 过捡，正样本判断为异常
        else:
            pred_thresh.append(0)
            if y_true[i] == 1:
                false_n.append(img_path_list[i])  # 漏检，异常判断为正常

    cm = confusion_matrix(y_true, pred_thresh)
    print(cm)
    print('false positive')
    print(false_p)
    print('false negative')
    print(false_n)

    FP = cm[0][1]
    TP = cm[0][0]
    FN = cm[1][0]
    TN = cm[1][1]

    noraml_recall = TP/(FP+TP)   # 正常召回率
    abnormal_recall = TN/(FN+TN)  # 异常召回率
    precision = (TN+TP)/(FP+TP+FN+TN)  # 正常异常准确率

    print("正常召回率:", noraml_recall)
    print("异常召回率:", abnormal_recall)
    print("正常异常准确率:", precision)

    return noraml_recall, abnormal_recall, precision, false_p, false_n


def estimate_thr_recall(y_true, y_pred_p, draw=False):
    from matplotlib import pyplot
    # 1. default 
    precision, recall, thresholds = precision_recall_curve(y_true, y_pred_p)

    # convert to f score
    fscore = (2 * precision * recall) / (precision + recall)
    # locate the index of the largest f score
    ix = np.argmax(fscore)
    print('Best Threshold=%f, F-Score=%.3f' % (thresholds[ix], fscore[ix]))

    if draw:
        abnormal = len(y_true[y_true==1]) / len(y_true)   # 异常比率
        pyplot.plot([0,1], [abnormal, abnormal], linestyle='--', label='abnormal')
        pyplot.plot(recall, precision, marker='.', label='Logistic')
        pyplot.scatter(recall[ix], precision[ix], marker='o', color='black', label='Best')
        # axis labels
        pyplot.xlabel('Recall')
        pyplot.ylabel('Precision')
        pyplot.legend()
        # show the plot
        # pyplot.show()
        pyplot.savefig('./roc.png')

    return thresholds[ix]


def sigmoid(x):
    # 改为向量版本
    return np.where(x >= 0, 1 / (1 + np.exp(-x)), np.exp(x) / (1 + np.exp(x)))

def to_labels(pos_probs, threshold):
	return (pos_probs >= threshold).astype('int')


def estimate_thr_by_list(y_true, y_pred_p, step_p=0.01):
    # 2. define thresholds
    y_pred_p = sigmoid(y_pred_p)  # may logistic
    thresholds = list(range(0, 1, step_p))
    # evaluate each threshold
    scores = [f1_score(y_true, to_labels(y_pred_p, t)) for t in thresholds]
    # get best threshold
    ix = np.argmax(scores)
    print('Threshold=%.3f, F-Score=%.5f' % (thresholds[ix], scores[ix]))

    return thresholds[ix]


# add from favae====='
import pandas as pd
from numpy import ndarray as NDArray
import skimage
import statistics
import sklearn
import tqdm


# __all__ = ['AverageMeter', 'compute_pro_score', 'estimate_thred_with_fpr']

class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


# pro_score 查看区别于semi_ort_mvtec
def compute_pro_score(amaps: NDArray, masks: NDArray) -> float:
    # from skimage import measure
    # from statistics import mean

    df = pd.DataFrame([], columns=["pro", "fpr", "threshold"])
    binary_amaps = np.zeros_like(amaps, dtype=np.bool)

    max_step = 200
    min_th = amaps.min()
    max_th = amaps.max()
    delta = (max_th - min_th) / max_step

    for th in tqdm(np.arange(min_th, max_th, delta), desc="compute pro"):
        binary_amaps[amaps <= th] = 0
        binary_amaps[amaps > th] = 1

        pros = []
        for binary_amap, mask in zip(binary_amaps, masks):
            for region in skimage.measure.regionprops(skimage.measure.label(mask)):
                axes0_ids = region.coords[:, 0]
                axes1_ids = region.coords[:, 1]
                TP_pixels = binary_amap[axes0_ids, axes1_ids].sum()
                pros.append(TP_pixels / region.area)

        inverse_masks = 1 - masks
        FP_pixels = np.logical_and(inverse_masks, binary_amaps).sum()
        fpr = FP_pixels / inverse_masks.sum()

        df = df.append({"pro": statistics.mean(pros), "fpr": fpr, "threshold": th}, ignore_index=True)

    return sklearn.metrics.auc(df["fpr"], df["pro"])


# 区别于 estimate_thr_recall
def estimate_thred_with_fpr(scores, max_step=1000, expect_fpr=0.05):
    """
    Use training or validation set to estimate the threshold.
    """
    threshold = 0
    min_th = scores.min()
    max_th = scores.max()
    delta = (max_th - min_th) / max_step
    for step in range(max_step):
        threshold = max_th - step * delta
        # segmentation
        binary_score_maps = np.zeros_like(scores)
        binary_score_maps[scores <= threshold] = 0
        binary_score_maps[scores > threshold] = 1

        # estimate the optimal threshold base on user defined min_area
        fpr = binary_score_maps.sum() / binary_score_maps.size

        if fpr >= expect_fpr:  # find the optimal threshold
            break
        
    return threshold



def write_results(args, img_rocauc, pixel_rocauc, noraml_recall, abnormal_recall, precision, best_f1_v,false_p, false_n, save_name):
    #  统一写性能文件

    values = dict()

    values = {'img_rocauc': round(img_rocauc, 4), 'pixel_rocauc': round(pixel_rocauc, 4)}

    values["noraml_recall"] = round(noraml_recall, 4)
    values["abnormal_recall"] = round(abnormal_recall, 4)
    values["precision"] = round(precision, 4)

    # add best thresh
    values["best_f1thr"] = round(best_f1_v, 6)

    # add pre wrong sample's file path
    values["false_p"] = false_p
    values["false_n"] = false_n

    if not os.path.exists(args.result_dir):
        os.makedirs(args.result_dir, exist_ok=True)

    with open(os.path.join(args.result_dir, save_name), 'a') as f:
        f.write(args.category + " : " + str(values) + "\n")